# trading/simulator.py
import math
from typing import Dict
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from data.mysql_db import DataManager


class TradingSimulator:
    def __init__(self, signals_df: pd.DataFrame, initial_cash=100000, commission_rate=0.0001):
        """
        初始化模拟交易模块

        参数:
            signals_df (DataFrame): 包含交易信号的 DataFrame
            initial_cash (float): 初始资金
            commission_rate (float): 佣金费率
        """
        self.signals_df = signals_df.copy()
        self.initial_cash = initial_cash
        self.commission_rate = commission_rate
        self._portfolio = None

    def execute_trades(self) -> pd.DataFrame:
        """执行交易并计算投资组合价值"""
        df = self.signals_df.copy()

        # 初始化账户（显式设置为 float 类型）
        df['Cash'] = self.initial_cash * 1.0
        df['Position'] = 0.0
        df['Portfolio_Value'] = self.initial_cash * 1.0

        cash = self.initial_cash
        position = 0.0
        portfolio_value = self.initial_cash
        trades = []

        for idx, row in df.iterrows():
            close_price = row['close']
            signal = row['Signal']

            # 执行交易
            if signal == 1 and position <= 0:  # 买入信号
                if close_price <= 0:
                    continue  # 价格为0时跳过

                # 计算可买入股数（向下取整为100的倍数）
                max_shares = math.floor(cash / close_price)
                shares = (max_shares // 100) * 100

                # 调整股数，确保佣金后总成本 <= 可用现金
                while shares > 0:
                    cost = shares * close_price * (1 + self.commission_rate)
                    if cost <= cash:
                        break
                    shares -= 100  # 向下调整100股

                if shares > 0:
                    cash -= cost
                    position += shares
                    trade_value = shares * close_price
                    trades.append({
                        'Date': idx,
                        'Action': 'Buy',
                        'Price': close_price,
                        'Shares': shares,
                        'Value': trade_value
                    })

            elif signal == -1 and position >= 0:  # 卖出信号
                if position > 0:
                    proceeds = position * close_price * (1 - self.commission_rate)
                    cash += proceeds
                    position = 0.0
                    trade_value = proceeds
                    trades.append({
                        'Date': idx,
                        'Action': 'Sell',
                        'Price': close_price,
                        'Shares': 0,
                        'Value': trade_value
                    })

            # 更新组合价值
            portfolio_value = cash + position * close_price
            df.loc[idx, 'Cash'] = cash
            df.loc[idx, 'Position'] = position
            df.loc[idx, 'Portfolio_Value'] = portfolio_value

        # 记录交易记录（处理空交易情况）
        if not trades:
            df['Trades'] = pd.NA  # 使用 pd.NA 替代 {}
        else:
            trade_dict = {t['Date']: t for t in trades}
            df['Trades'] = df.index.map(trade_dict).fillna(pd.NA)  # 使用 pd.NA 替代 {}

        self._portfolio = df
        return df

    def calculate_performance(self) -> Dict:
        """计算绩效指标"""
        df = self._portfolio

        # 基本绩效指标
        total_return = (df['Portfolio_Value'].iloc[-1] / self.initial_cash - 1) * 100
        annualized_return = (df['Portfolio_Value'].iloc[-1] / self.initial_cash) ** (252 / len(df)) - 1
        max_drawdown = self._calculate_max_drawdown(df['Portfolio_Value'])
        sharpe_ratio = self._calculate_sharpe_ratio(df['Portfolio_Value'].pct_change().dropna())

        # 交易统计
        total_trades = 0
        winning_trades = 0

        for t in df['Trades']:
            if isinstance(t, dict):  # 确保交易记录是字典
                if t['Action'] == 'Sell':  # 仅统计卖出交易
                    total_trades += 1

                    # 获取前一日收盘价
                    prev_date = t['Date'] - pd.Timedelta(days=1)
                    if prev_date in df.index:
                        prev_close = df.loc[prev_date, 'close']
                        if t['Value'] > prev_close * t['Shares']:
                            winning_trades += 1

        return {
            'Final_Value': df['Portfolio_Value'].iloc[-1],
            'Total_Return': total_return,
            'Annualized_Return': annualized_return,
            'Max_Drawdown': max_drawdown,
            'Sharpe_Ratio': sharpe_ratio,
            'Total_Trades': total_trades,
            'Winning_Trades': winning_trades
        }

    def _calculate_max_drawdown(self, portfolio_value):
        """计算最大回撤"""
        peak = portfolio_value.cummax()
        drawdown = (portfolio_value - peak) / peak
        return drawdown.min() * 100

    def _calculate_sharpe_ratio(self, returns, risk_free_rate=0.02):
        """计算夏普比率"""
        if len(returns) < 2:
            return 0
        excess_returns = returns - risk_free_rate / 252
        return np.sqrt(252) * excess_returns.mean() / returns.std()

    def plot_results(self):
        """绘制净值曲线和交易信号"""
        df = self._portfolio

        plt.figure(figsize=(14, 7))

        # 绘制净值曲线
        plt.subplot(2, 1, 1)
        df['Portfolio_Value'].plot(label='Portfolio Value')
        df['Portfolio_Value'].rolling(window=20).mean().plot(label='20-Day MA')
        plt.title('Portfolio Value Over Time')
        plt.legend()

        # 绘制交易信号
        plt.subplot(2, 1, 2)
        df['close'].plot(label='Close Price')
        df['DIF'].plot(label='MACD Line')
        df['MACD'].plot(label='Signal Line')

        # 标记买入信号
        buy_idx = df[df['Signal'] == 1].index
        plt.scatter(buy_idx, df.loc[buy_idx, 'close'], color='green', marker='^', alpha=1, label='Buy Signal')

        # 标记卖出信号
        sell_idx = df[df['Signal'] == -1].index
        plt.scatter(sell_idx, df.loc[sell_idx, 'close'], color='red', marker='v', alpha=1, label='Sell Signal')

        plt.title('Trading Signals')
        plt.legend()
        plt.tight_layout()
        plt.show()


# 在示例用法中添加调试日志
if __name__ == '__main__':
    from strategies.ma_strategy import MovingAverageStrategy
    from indicators.indicator_calculator import TechnicalIndicators

    ts_code = '600610.SH'
    data_manager = DataManager()
    df = data_manager.get_stock_daily_data(ts_code)
    indicator = TechnicalIndicators(df)
    strategy = MovingAverageStrategy(indicator=indicator, short_window=12, long_window=26)

    # 生成交易信号
    signals_df = strategy.generate_signals()
    print("交易信号统计:")
    print(signals_df['Signal'].value_counts())

    # 模拟交易
    simulator = TradingSimulator(signals_df=signals_df, initial_cash=100000)
    result_df = simulator.execute_trades()
    print("交易记录示例:")
    print(result_df[['trade_date', 'Signal', 'Trades']].head())
    result_df.to_csv('simulation_result.csv')
    # 计算绩效
    performance = simulator.calculate_performance()
    print("回测结果:")
    print(f"最终资产: {performance['Final_Value']:.2f}")
    print(f"总收益率: {performance['Total_Return']:.2f}%")
    print(f"年化收益率: {performance['Annualized_Return']:.2f}%")
    print(f"最大回撤: {performance['Max_Drawdown']:.2f}%")
    print(f"夏普比率: {performance['Sharpe_Ratio']:.2f}")
    print(f"总交易次数: {performance['Total_Trades']}")
    print(f"盈利交易次数: {performance['Winning_Trades']}")

    # 可视化结果
    simulator.plot_results()
